# -*- coding: utf-8 -*-
"""
Created on Wed Sep 23 10:41:00 2020

@author: Robin
"""
from .utils import process_raw_data
#from .IOFile import readjson
from DataIO import pymysqlops
#import datetime
import numpy as np
import pandas as pd

cla_stock_mkt = pymysqlops('127.0.0.1',3306,'root','admin123','stockInfo')
cla_stock_mkt.db_connection()

readpath = 'data'


'''
参数不匹配
'''
class NonMatchingParaError(Exception):
    pass

'''
根据传入参数来获取资产数据
''' 
def cal_stockdf(stock_codes,begdate,enddate,
                 period=1,
                 include_st = False,
                 include_suspend = False,
                 include_new_stock = False,
                 ipo_days = 60):         
    '''
    begdate = '2020-08-01'
    enddate = '2020-09-01'
    stock_codes = ['000008.SZ','000012.SZ','600566.SH','600567.SH']
    '''
    if not isinstance(period,int):
        raise NonMatchingParaError("period should be integate")
    
    # 提取相关代码的数据    
    if isinstance(stock_codes,list):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "'and TRADEDATE<= '" + enddate + "'"    
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "'"  
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE <= '"+\
                      enddate + "'"                        
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt"
       
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
        df_stock_clear['close'] = pd.DataFrame(df_stock_clear['close'] ,dtype = np.float)     
        # 筛选相关股票代码的数据
        df_stock_clear = df_stock_clear[df_stock_clear['asset'].isin(stock_codes)]
    
    elif isinstance(stock_codes,str):
        # SQL命令：直接筛选个股数据
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "'and TRADEDATE<= '" + enddate + "' and code = '" + stock_codes +"'" 
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "' and code = '" + stock_codes +"'"  
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE <= '"+\
                      enddate + "' and code = '" + stock_codes +"'"                        
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where code = '" + stock_codes +"'" 
                               
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
        df_stock_clear['close'] = pd.DataFrame(df_stock_clear['close'] ,dtype = np.float)        

    elif isinstance(stock_codes,pd.MultiIndex):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "'and TRADEDATE<= '" + enddate + "'"    
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE >= '"+\
                      begdate + "'"  
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt where TRADEDATE <= '"+\
                      enddate + "'"                        
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code,close FROM stockinfo.stock_mkt"
                          
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
        df_stock_clear['close'] = pd.DataFrame(df_stock_clear['close'] ,dtype = np.float)     
        # 筛选相关股票代码的数据
        df_stock_clear = pd.merge(df_stock_clear,stock_codes.to_frame(index= False),on = ['date','asset'])

    # 提取步长要求的日期
    if period != 1:
        date_series = df_stock_clear['date'].unique()
        date_series = date_series[0:-1:period]       
        df_stock_clear = df_stock_clear[df_stock_clear['date'].isin(date_series)]
        
    # 是否剔除st股票
    if not include_st:
        pass
    
    # 是否剔除当天停牌的股票
    if not include_suspend:
        pass

    # 是否剔除新上市的股票
    if not include_new_stock:
        pass
    
    # 格式调整
    df_stock_clear['date'] = df_stock_clear['date'].apply(pd.to_datetime)        
    df_stock_clear.set_index(['date','asset'],inplace = True)
    
    return df_stock_clear


def get_indi_data(stock_codes, ind_code,begdate=None,enddate=None,
             period=1):
    '''
    提取单因子数据，
    表名称为ind_+因子名称,
    字段名称为 因子名称
    '''
    '''
    begdate = '2020-08-01'
    enddate = '2020-09-01'
    stock_codes = ['000008.SZ','000012.SZ','600566.SH','600567.SH']
    ind_codes = "weight"
    '''

    if not isinstance(period,int):
        raise NonMatchingParaError("period should be integate")    

    # SQL命令
    if (begdate is not None) & (enddate is not None):
        str_sql = "select tradedate,code," + ind_code + " FROM stockinfo.ind_" + \
                    ind_code +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "'"
    elif (begdate is None) & (enddate is not None ):
        str_sql = "select tradedate,code," + ind_code + " FROM stockinfo.ind_" + \
                    ind_code +  " where TRADEDATE<= '" + enddate + "'"  
    elif (begdate is not None) & (enddate is None ):
        str_sql = "select tradedate,code," + ind_code + " FROM stockinfo.ind_" + \
                    ind_code +  " where TRADEDATE>= '" + begdate + "'"  
    elif (begdate is None) & (enddate is not None ):
        str_sql = "select tradedate,code," + ind_code + " FROM stockinfo.ind_" + ind_code

    # 数据提取
    result = cla_stock_mkt.select_table_by_sql(str_sql)
    df_stock = pd.DataFrame(list(result)).rename(columns = {0:'date',1:'asset',2:'close'})
    df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
    df_stock_clear['close'] = pd.DataFrame(df_stock_clear['close'] ,dtype = np.float)     
    # 筛选相关股票代码的数据
    df_stock_clear = df_stock_clear[df_stock_clear['asset'].isin(stock_codes)]
    
    # 提取步长要求的日期
    if period != 1:
        date_series = df_stock_clear['date'].unique()
        date_series = date_series[0:-1:period]       
        df_stock_clear = df_stock_clear[df_stock_clear['date'].isin(date_series)] 

    # 格式调整
    df_stock_clear['date'] = df_stock_clear['date'].apply(pd.to_datetime)        
    df_stock_clear.set_index(['date','asset'],inplace = True)
    
    return df_stock_clear

'''
根据传入参数来获取因子的数据
'''      
def cal_inddf(stock_codes,ind_codes,begdate=None,enddate=None,period=1):  
    '''
    begdate = '2020-08-01'
    enddate = '2020-09-01'
    stock_codes = ['000008.SZ','000012.SZ','600566.SH','600567.SH']
    ind_codes = "weight"
    '''
    # 只提取一个因子
    if isinstance(ind_codes,str):
        if isinstance(stock_codes,pd.MultiIndex):    
            stock_code_list = stock_codes.index.levels[1]
            factor = get_indi_data(stock_code_list, ind_codes, begdate, enddate, period)     # 提取数据
            factors = pd.merge(factor,stock_codes.to_frame(index= False),on = ['date','asset']) 
        else:
            factor = get_indi_data(stock_codes, ind_codes, begdate, enddate, period)     # 提取数据
            factors = factor[factor['asset'].isin(stock_codes)]        
        
    # 提取多个因子
    elif isinstance(ind_codes,list):     
        if isinstance(stock_codes,pd.MultiIndex):    
            factors = stock_codes.to_frame(index= False)
            for ind_code in ind_codes:
                stock_code_list = stock_codes.index.levels[1]
                factor = get_indi_data(stock_code_list, ind_codes, begdate, enddate, period)     # 提取数据
                factors = pd.merge(factors,factor,on = ['date','asset']) 
        else:
            factors = {}
            for ind_code in ind_codes:            
                factor = get_indi_data(stock_codes, ind_codes, begdate, enddate, period)     # 提取数据
                factor = factor[factor['asset'].isin(stock_codes)]   
                try:
                    factors = pd.merge(factors,factor,on = ['date','asset']) 
                except:
                    factors = factor.copy(deep = True)
        
    #调用utils.py里的process_raw_data函数进行数据标准化
    factors=process_raw_data(factors)    
    return factors


'''
行业数据
'''
def cal_groupby(stock_codes,
                begdate=None,
                enddate=None,  
                period = 1,
                industry_type='sw_1'):
    '''
    groupby : pd.Series - MultiIndex or dict
                Either A MultiIndex Series indexed by date and asset,
                containing the period wise group codes for each asset, or
                a dict of asset to group mappings. If a dict is passed,
                it is assumed that group mappings are unchanged for the
                entire time period of the passed factor data.
    '''     
    '''
    建议传入数据为multiindex
    '''
    '''
    提取单行业划分数据，
    表名称为indust_ _+ 划分依据 + 级别 （申万一级）,
    字段名称为 划分依据 + 级别 （申万一级）
    '''
    '''
    industry_type='sw_1'
    begdate = '1990-12-10'
    enddate = '1992-12-31'
    period= 8
    stock_codes = ['000001.SZ','000008.SZ','000012.SZ','600566.SH','600567.SH']    
    '''    
    if not isinstance(period,int):
        raise NonMatchingParaError("period should be integate")    
        
    # 提取相关代码的数据    
    if isinstance(stock_codes,list):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "'"  
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate +  "'"   
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE<= '" + enddate + "'"                       
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type
               
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')  
        # 筛选相关股票代码的数据
        df_stock_clear = df_stock_clear[df_stock_clear['asset'].isin(stock_codes)]    
    elif isinstance(stock_codes,str):
        # SQL命令：直接筛选个股数据
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "' and code = '" + stock_codes +"'" 
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate +  "' and code = '" + stock_codes +"'"  
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE<= '" + enddate + "' and code = '" + stock_codes +"'"                        
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type + " where code = '" + stock_codes +"'"         
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
    elif isinstance(stock_codes,pd.MultiIndex):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "'"  
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE >= '"+begdate +  "'"   
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type +  " where TRADEDATE<= '" + enddate + "'"                       
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + industry_type + " FROM stockinfo.indust_" + \
                        industry_type               
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')   
        # 筛选相关股票代码的数据
        df_stock_clear = pd.merge_asof(stock_codes.to_frame(index= False),df_stock_clear,on = ['date','asset'])
        
    return df_stock_clear

'''
个股权重数据
'''
def cal_weightdf(stock_codes,
                 base_code='000905.SH',
                 begdate='1900-01-01',
                 enddate='2099-12-31',
                 period = 1):
    
    if not isinstance(period,int):
        raise NonMatchingParaError("period should be integate")    
        
    # 提取相关代码的数据    
    if isinstance(stock_codes,list):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "'"  
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate +  "'"   
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE<= '" + enddate + "'"                       
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code
                                     
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')  
        # 筛选相关股票代码的数据
        df_stock_clear = df_stock_clear[df_stock_clear['asset'].isin(stock_codes)]    
    elif isinstance(stock_codes,str):
        # SQL命令：直接筛选个股数据
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "' and code = '" + stock_codes +"'"  
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate +  "' and code = '" + stock_codes +"'"  
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE<= '" + enddate + "' and code = '" + stock_codes +"'"                        
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code + " where code = '" + stock_codes +"'" 
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')
    elif isinstance(stock_codes,pd.MultiIndex):
        # SQL命令
        if (begdate is not None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate + "'and TRADEDATE<= '" + enddate + "'"  
        elif (begdate is not None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE >= '"+begdate +  "'"   
        elif (begdate is None) & (enddate is not None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code +  " where TRADEDATE<= '" + enddate + "'"                       
        elif (begdate is None) & (enddate is None):
            str_sql = "select tradedate,code," + base_code + " FROM stockinfo.indust_" + \
                        base_code               
        # 数据提取
        stock_data = cla_stock_mkt.select_table_by_sql(str_sql)
        df_stock = pd.DataFrame(list(stock_data)).rename(columns = {0:'date',1:'asset',2:'close'})
        df_stock_clear = df_stock.dropna(axis = 'index', how = 'any')   
        # 筛选相关股票代码的数据
        df_stock_clear = pd.merge_asof(stock_codes.to_frame(index= False),df_stock_clear,on = ['date','asset'])
        
    return df_stock_clear


